from langchain_community.llms import Tongyi

# 设置查询问题
# query = "客户经理被投诉了，投诉一次扣多少分？"
query = "客户经理每年评聘申报时间是怎样的？"
if query:
    # 示例：如何加载已保存的向量数据库
    # 注释掉以下代码以避免在当前运行中重复加载
    # 创建嵌入模型
    embeddings = DashScopeEmbeddings(
        model="text-embedding-v2"
    )
    # 从磁盘加载向量数据库
    loaded_knowledgeBase = load_knowledge_base("./vector_db", embeddings)
    # 使用加载的知识库进行查询
    docs = loaded_knowledgeBase.similarity_search(query)

    # 初始化对话大模型
    DASHSCOPE_API_KEY = os.getenv("DASHSCOPE_API_KEY"),
    llm = Tongyi(model_name="deepseek-v3", dashscope_api_key=DASHSCOPE_API_KEY)

    # 加载问答链
    chain = load_qa_chain(llm, chain_type="stuff")

    # 准备输入数据
    input_data = {"input_documents": docs, "question": query}

    # 使用回调函数跟踪API调用成本
    with get_openai_callback() as cost:
        # 执行问答链
        response = chain.invoke(input=input_data)
        print(f"查询已处理。成本: {cost}")
        print(response["output_text"])
        print("来源:")

    # 记录唯一的页码
    unique_pages = set()

    # 显示每个文档块的来源页码
    for doc in docs:
        text_content = getattr(doc, "page_content", "")
        source_page = knowledgeBase.page_info.get(
            text_content.strip(), "未知"
        )

        if source_page not in unique_pages:
            unique_pages.add(source_page)
            print(f"文本块页码: {source_page}")